import tkinter as tk
from tkinter import ttk, scrolledtext
import threading
import time
import re
import os
import json
from openai import OpenAI
from encryption_utils import EncryptionUtils

class SidebarAIAssistant:
    def __init__(self, parent, terminal_app):
        self.parent = parent
        self.terminal_app = terminal_app
        self.is_visible = False
        self.current_tab_index = -1
        self.chat_histories = {}
        self.assistant_thread = None
        self.stop_event = threading.Event()
        
        # 初始化加密工具并获取API密钥
        self.encryption_utils = EncryptionUtils()
        api_key = self.get_api_key()
        
        # 保存当前使用的密钥名称
        self.current_key_name = self._load_selected_key_name()
        
        # 初始化DeepSeek客户端
        self.client = OpenAI(
            api_key=api_key,
            base_url="https://api.deepseek.com"
        )
        
        # 创建侧边栏UI
        self.create_ui()
        
        # 绑定标签页切换事件
        self.terminal_app.tab_control.bind("<<NotebookTabChanged>>", self.on_tab_changed)
    
    def get_api_key(self):
        """从密钥管理器获取API密钥"""
        try:
            keys = self.encryption_utils.load_keys()
            # 如果没有密钥，返回空字符串
            if not keys:
                return ""
            
            # 尝试加载选定的密钥
            selected_key_name = self._load_selected_key_name()
            if selected_key_name and selected_key_name in keys:
                return keys[selected_key_name]
            
            # 如果没有找到选定的密钥，但有默认密钥名称，尝试使用它
            if hasattr(self, 'current_key_name') and self.current_key_name and self.current_key_name in keys:
                return keys[self.current_key_name]
            
            # 查找DeepSeek相关的密钥
            for name, key in keys.items():
                if "deepseek" in name.lower():
                    return key
            
            # 如果没有找到特定名称的密钥，返回第一个密钥
            return list(keys.values())[0]
        except Exception as e:
            print(f"获取API密钥时出错: {e}")
            return ""
    
    def _load_selected_key_name(self):
        """加载选定的密钥名称"""
        try:
            if os.path.exists("selected_key.json"):
                with open("selected_key.json", "r", encoding="utf-8") as f:
                    data = json.load(f)
                    return data.get("selected_key", "")
            return ""
        except Exception:
            return ""
    
    def create_ui(self):
        # 创建侧边栏框架 (约6cm宽度)
        self.sidebar_frame = ttk.Frame(self.terminal_app.paned_window, width=180, height=600, relief=tk.RAISED, borderwidth=1)
        self.sidebar_frame.pack_propagate(False)
        # 设置最小尺寸
        self.sidebar_frame.columnconfigure(0, minsize=180)
        
        # 创建标题栏
        title_frame = ttk.Frame(self.sidebar_frame)
        title_frame.pack(fill=tk.X, padx=5, pady=5)
        
        title_label = ttk.Label(title_frame, text="侧边栏AI助手", font=('SimHei', 12, 'bold'))
        title_label.pack(side=tk.LEFT, padx=5)
        
        # 增大关闭按钮的尺寸，使其更容易点击
        close_button = ttk.Button(title_frame, text="×", width=5, style="Accent.TButton", command=self.toggle_visibility)
        close_button.pack(side=tk.RIGHT, padx=5)
        
        # 创建聊天显示区域
        self.chat_display = scrolledtext.ScrolledText(
            self.sidebar_frame, wrap=tk.WORD, font=('SimHei', 10), state=tk.DISABLED, bg='#f0f0f0'
        )
        self.chat_display.pack(fill=tk.BOTH, expand=True, padx=2, pady=(0, 2))
        
        # 创建输入区域
        input_frame = ttk.Frame(self.sidebar_frame)
        input_frame.pack(fill=tk.X, padx=2, pady=2)
        # print("tk.x:",tk.X)
        
        self.input_entry = scrolledtext.ScrolledText(input_frame, wrap=tk.WORD, font=('SimHei', 10), height=3, borderwidth=1, relief=tk.SUNKEN)
        self.input_entry.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=(0, 5))
        self.input_entry.bind("<Return>", self.handle_input_return)
        
        # 创建按钮框架，确保按钮始终可见
        button_frame = ttk.Frame(input_frame)
        button_frame.pack(side=tk.RIGHT, fill=tk.Y)
        
        send_button = ttk.Button(button_frame, text="发送", width=6, command=self.send_message)
        send_button.pack(side=tk.TOP, pady=(0, 2))
        
        clear_button = ttk.Button(button_frame, text="清空", width=6, command=self.clear_chat)
        clear_button.pack(side=tk.TOP)
        
        # 初始化气泡效果的样式
        self.chat_display.tag_config("user_name", foreground="#0000ff", font=('SimHei', 10, 'bold'))
        self.chat_display.tag_config("assistant_name", foreground="#ff0000", font=('SimHei', 10, 'bold'))
        self.chat_display.tag_config("typing_tag", foreground="gray", font=('SimHei', 10, 'italic'))
        
        # 创建ttk样式
        self.style = ttk.Style()
        self.style.configure("User.TLabel", background="#dcf8c6", foreground="#000000", font=('SimHei', 10))
        self.style.configure("Assistant.TLabel", background="#ffffff", foreground="#000000", font=('SimHei', 10))
        # 定义关闭按钮的样式
        self.style.configure("Accent.TButton", foreground="#ff0000", font=('SimHei', 10, 'bold'))
    
    def toggle_visibility(self):
        print(f"Toggle visibility called, current state: {self.is_visible}")
        
        # 获取当前标签页索引
        current_tab_idx = self.terminal_app.tab_control.index(self.terminal_app.tab_control.select())
        self.current_tab_index = current_tab_idx
        
        if self.is_visible:
            # 隐藏侧边栏 - 完全从PanedWindow中移除
            try:
                # 从PanedWindow中彻底移除侧边栏框架
                self.terminal_app.paned_window.remove(self.sidebar_frame)
                print("隐藏侧边栏成功")
            except Exception as e:
                print(f"隐藏侧边栏时出错: {str(e)}")
        else:
            # 显示侧边栏 - 直接添加到PanedWindow
            try:
                # 检查侧边栏是否已经在PanedWindow中，如果不在则添加
                try:
                    self.terminal_app.paned_window.add(self.sidebar_frame, weight=2)
                except Exception:
                    # 如果已经存在，则先移除再添加，确保正确显示
                    try:
                        self.terminal_app.paned_window.remove(self.sidebar_frame)
                    except Exception:
                        pass
                    self.terminal_app.paned_window.add(self.sidebar_frame, weight=2)
                
                # 确保侧边栏框架可见，但保持尺寸控制
                self.sidebar_frame.pack_propagate(False)
                
                # 如果当前标签页没有聊天历史，初始化
                if current_tab_idx not in self.chat_histories:
                    self.chat_histories[current_tab_idx] = []
                    # 显示欢迎消息
                    self.display_message("助手", "你好！我是你的侧边栏AI助手，有什么可以帮助你的吗？")
                else:
                    # 加载当前标签页的聊天历史
                    self.load_chat_history(current_tab_idx)
                
                # 自动聚焦到输入框
                self.input_entry.focus_set()
                
                print("显示侧边栏成功")
            except Exception as e:
                print(f"显示侧边栏时出错: {str(e)}")
                import traceback
                traceback.print_exc()
        
        # 更新可见性状态
        self.is_visible = not self.is_visible
    
    def on_tab_changed(self, event):
        # 切换标签页时，如果侧边栏可见，更新聊天历史
        if self.is_visible:
            # 保存当前聊天历史
            if self.current_tab_index != -1:
                # 聊天历史已经通过display_message实时保存，这里不需要额外保存
                pass
            
            # 更新当前标签页索引
            current_tab_idx = self.terminal_app.tab_control.index(self.terminal_app.tab_control.select())
            self.current_tab_index = current_tab_idx
            
            # 如果当前标签页没有聊天历史，初始化
            if current_tab_idx not in self.chat_histories:
                self.chat_histories[current_tab_idx] = []
                # 显示欢迎消息
                self.display_message("助手", "你好！我是你的侧边栏AI助手，有什么可以帮助你的吗？")
            else:
                # 加载当前标签页的聊天历史
                self.load_chat_history(current_tab_idx)
        else:
            # 侧边栏不可见时也要更新当前标签页索引，确保后续操作正确
            current_tab_idx = self.terminal_app.tab_control.index(self.terminal_app.tab_control.select())
            self.current_tab_index = current_tab_idx
            # 如果当前标签页没有聊天历史，初始化
            if current_tab_idx not in self.chat_histories:
                self.chat_histories[current_tab_idx] = []
    
    def load_chat_history(self, tab_idx):
        # 清空聊天显示区域
        self.chat_display.config(state=tk.NORMAL)
        self.chat_display.delete(1.0, tk.END)
        self.chat_display.config(state=tk.DISABLED)
        
        # 批量加载聊天历史，避免频繁更新UI
        self.chat_display.config(state=tk.NORMAL)
        for i, (role, message) in enumerate(self.chat_histories[tab_idx]):
            # 添加角色标签和气泡
            if role == "你":
                # 用户消息 - 右对齐气泡
                self.chat_display.insert(tk.END, f"{role}\n", "user_name")
                
                # 插入用户气泡
                self.chat_display.window_create(tk.END, window=self.create_message_bubble(message, is_user=True))
            else:
                # 助手消息 - 左对齐气泡
                self.chat_display.insert(tk.END, f"{role}\n", "assistant_name")
                
                # 插入助手气泡
                self.chat_display.window_create(tk.END, window=self.create_message_bubble(message, is_user=False))
            
            # 添加空行分隔消息
            self.chat_display.insert(tk.END, "\n\n")
            
            # 每加载10条消息后稍作停顿，避免界面卡死
            if i % 10 == 9:
                self.chat_display.config(state=tk.DISABLED)
                self.parent.update_idletasks()
                self.chat_display.config(state=tk.NORMAL)
        
        self.chat_display.config(state=tk.DISABLED)
        self.chat_display.see(tk.END)
    
    def display_message(self, role, message):
        # 保存到聊天历史
        if self.current_tab_index != -1:
            self.chat_histories[self.current_tab_index].append((role, message))
        
        # 显示消息
        self.chat_display.config(state=tk.NORMAL)
        
        # 添加角色标签和气泡
        if role == "你":
            # 用户消息 - 右对齐气泡
            self.chat_display.insert(tk.END, f"{role}\n", "user_name")
            
            # 插入用户气泡
            bubble_frame = self.create_message_bubble(message, is_user=True)
            self.chat_display.window_create(tk.END, window=bubble_frame)
        else:
            # 助手消息 - 左对齐气泡
            self.chat_display.insert(tk.END, f"{role}\n", "assistant_name")
            
            # 插入助手气泡
            bubble_frame = self.create_message_bubble(message, is_user=False)
            self.chat_display.window_create(tk.END, window=bubble_frame)
        
        # 添加空行分隔消息
        self.chat_display.insert(tk.END, "\n\n")
        
        self.chat_display.config(state=tk.DISABLED)
        self.chat_display.see(tk.END)
        
        # 强制刷新界面，确保消息立即显示
        self.parent.update_idletasks()
    
    def create_message_bubble(self, message, is_user=False):
        # 创建气泡框架
        bubble_frame = tk.Frame(self.chat_display, bg=self.chat_display['bg'])
        
        # 根据是否为用户消息选择不同的样式
        style_name = "User.TLabel" if is_user else "Assistant.TLabel"
        
        # 过滤掉Tkinter不支持的Unicode字符（超过U+FFFF范围的字符）
        filtered_message = ''.join(c for c in message if ord(c) <= 0xFFFF)
        
        # 创建消息标签
        message_label = ttk.Label(
            bubble_frame, 
            text=filtered_message, 
            wraplength=150, 
            justify=tk.LEFT,
            padding=(5, 3),
            style=style_name
        )
        
        # 根据用户类型设置对齐方式
        if is_user:
            message_label.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True)
        else:
            message_label.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        
        return bubble_frame
    
    def handle_input_return(self, event):
        # 检查是否按下了Shift+Enter，如果是则插入换行
        if event.state & 0x1:  # Shift键被按下
            return  # 允许默认的换行行为
        else:
            # 否则发送消息
            self.send_message()
            return "break"  # 阻止默认行为
    
    def send_message(self):
        user_message = self.input_entry.get(1.0, tk.END).strip()
        if not user_message or self.current_tab_index == -1:
            return
        
        # 清空输入框
        self.input_entry.delete(1.0, tk.END)
        
        # 显示用户消息
        self.display_message("你", user_message)
        
        # 获取当前终端的历史信息
        terminal_history = self.get_current_terminal_history()
        
        # 在单独线程中处理AI响应，避免界面卡顿
        self.stop_event.clear()
        self.assistant_thread = threading.Thread(
            target=self.generate_assistant_response,
            args=(user_message, terminal_history)
        )
        self.assistant_thread.daemon = True
        self.assistant_thread.start()
    
    def get_current_terminal_history(self):
        # 获取当前终端标签页的命令历史
        if 0 <= self.current_tab_index < len(self.terminal_app.terminal_tabs):
            current_terminal = self.terminal_app.terminal_tabs[self.current_tab_index]
            # 获取命令历史
            command_history = current_terminal.command_history
            
            # 获取终端的完整文本内容作为上下文
            terminal_text = current_terminal.terminal_text.get(1.0, tk.END)
            
            return {
                "command_history": command_history,
                "terminal_text": terminal_text
            }
        return {}
    
    def update_api_key(self):
        """更新API密钥"""
        # 获取新的API密钥
        new_key_name = self._load_selected_key_name()
        new_api_key = self.get_api_key()
        
        # 检查密钥是否发生变化
        if new_key_name != self.current_key_name or new_api_key != self.client.api_key:
            # 更新当前密钥名称
            self.current_key_name = new_key_name
            
            # 更新客户端的API密钥
            self.client = OpenAI(
                api_key=new_api_key,
                base_url="https://api.deepseek.com"
            )
            
            print(f"侧边栏AI助手API密钥已更新，当前使用密钥: {new_key_name if new_key_name else '默认密钥'}")

    def generate_assistant_response(self, user_input, terminal_history="", terminal_output=""):
        # 显示"正在输入"状态
        self.parent.after(0, self.update_typing_status, True)
        
        try:
            # 更新API密钥以确保使用最新的密钥
            self.update_api_key()
            
            # 获取当前标签页的聊天历史
            messages = [
                {"role": "system", "content": "你是一个智能终端助手，帮助用户理解和使用命令行。你可以解释命令、提供建议并帮助用户解决问题，尽量不要长篇大论，，不要使用markdown格式，我的文本解析器不能显示markddown，如果你返回markdown会显得非常乱。内容从简，清晰明了，注意这是windows平台。"}
            ]
            
            # 添加聊天历史
            if self.current_tab_index != -1 and self.current_tab_index in self.chat_histories:
                for role, message in self.chat_histories[self.current_tab_index]:
                    messages.append({"role": "user" if role == "你" else "assistant", "content": message})
            
            # 如果有终端历史，添加到系统消息中
            if terminal_history:
                system_context = "\n\n用户最近在终端中执行的命令: "
                if terminal_history.get("command_history"):
                    system_context += "; ".join(terminal_history['command_history'][-5:])  # 只取最近5条命令
                
                # 添加终端的完整输出内容（限制在最后2000个字符以内以避免过长）
                if terminal_history.get("terminal_text"):
                    terminal_content = terminal_history['terminal_text'][-2000:]  # 只取最后2000个字符
                    system_context += f"\n\n终端最近的输出内容:\n{terminal_content}"
                
                messages[0]["content"] += system_context
            
            # 调用DeepSeek API
            response = self.client.chat.completions.create(
                model="deepseek-chat",
                messages=messages,
                stream=False
            )
            
            # 获取AI回复
            ai_message = response.choices[0].message.content.strip()
            
            # 检查是否需要停止
            if self.stop_event.is_set():
                return
            
            # 显示AI响应
            self.parent.after(0, self.update_typing_status, False)
            self.parent.after(0, self.display_message, "助手", ai_message)
        except Exception as e:
            self.parent.after(0, self.update_typing_status, False)
            self.parent.after(0, self.display_message, "助手", f"生成响应时出错: {str(e)}")
    
    def update_typing_status(self, is_typing):
        # 更新UI以显示或隐藏"正在输入"状态
        if is_typing:
            # 检查是否已经有"正在输入"提示
            current_text = self.chat_display.get(1.0, tk.END)
            if "正在输入...\n" not in current_text:
                self.chat_display.config(state=tk.NORMAL)
                self.chat_display.insert(tk.END, "正在输入...\n", "typing_tag")
                self.chat_display.config(state=tk.DISABLED)
                self.chat_display.see(tk.END)
        else:
            # 移除"正在输入"提示
            self.chat_display.config(state=tk.NORMAL)
            # 查找并删除"正在输入...\n"
            try:
                typing_start = self.chat_display.search("正在输入...\n", 1.0, tk.END)
                if typing_start:
                    typing_end = self.chat_display.index(f"{typing_start}+1line")
                    self.chat_display.delete(typing_start, typing_end)
            except:
                pass
            self.chat_display.config(state=tk.DISABLED)
    
    def clear_chat(self):
        # 清空当前标签页的聊天历史
        if self.current_tab_index != -1:
            self.chat_histories[self.current_tab_index] = []
            self.chat_display.config(state=tk.NORMAL)
            self.chat_display.delete(1.0, tk.END)
            self.chat_display.config(state=tk.DISABLED)
            # 显示欢迎消息
            self.display_message("助手", "你好！我是你的侧边栏AI助手，有什么可以帮助你的吗？")
    
    def stop(self):
        # 停止AI助手
        self.stop_event.set()
        if self.assistant_thread and self.assistant_thread.is_alive():
            self.assistant_thread.join(timeout=1)